import cv2
import numpy as np


def calculate_focal_length(known_width_mm, known_distance_mm, pixel_width):
    """纯数学计算焦距"""
    return (pixel_width * known_distance_mm) / known_width_mm


def detect_object_width(frame):
    """从摄像头帧中检测物体宽度"""
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blurred, 50, 150)

    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None, None

    max_contour = max(contours, key=cv2.contourArea)
    rect = cv2.minAreaRect(max_contour)
    box = cv2.boxPoints(rect).astype(np.int32)
    return max(rect[1]), box  # 返回像素宽度和检测框坐标


def realtime_focal_calibration(known_width_mm, known_distance_mm):
    """摄像头实时焦距标定"""
    cap = cv2.VideoCapture(1)
    if not cap.isOpened():
        raise RuntimeError("无法打开摄像头")

    print(f"请将 {known_width_mm}mm 宽的物体放在 {known_distance_mm}mm 距离处")
    print("按空格键捕获当前帧计算焦距，ESC键退出")

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # 实时检测物体
        width_px, box = detect_object_width(frame)

        # 显示检测结果
        if width_px is not None:
            cv2.drawContours(frame, [box], 0, (0, 255, 0), 2)
            cv2.putText(frame, f"Detected: {width_px:.1f}px",
                        (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

        cv2.imshow("Focal Calibration", frame)

        key = cv2.waitKey(1)
        if key == 32:  # 空格键捕获
            if width_px:
                focal_px = calculate_focal_length(known_width_mm, known_distance_mm, width_px)
                print(f"像素宽度: {width_px:.1f}px -> 计算焦距: {focal_px:.1f}px")
                cv2.putText(frame, f"Focal: {focal_px:.1f}px",
                            (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                cv2.imshow("Captured Result", frame)
        elif key == 27:  # ESC键退出
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    # 示例：使用A4纸（210mm宽）在1000mm距离标定
    realtime_focal_calibration(
        known_width_mm=210.0,
        known_distance_mm=1410.5 # 99 627  110.5 618.4    131.5 617.8   141.5 617.9
    )